from abc import ABC

import tensorflow as tf
from losses import ArcMarginProduct, CosMarginProduct


class FaceModel(tf.keras.Model, ABC):
    def __init__(self, dim=128, n_cls=1000, scale=30, margin=0.3, loss_type='arc'):
        super(FaceModel, self).__init__()
        assert loss_type in ("arc", "cos")

        self.backbone = tf.keras.applications.ResNet50(
            include_top=False,
            pooling="avg",
            input_shape=(224, 224, 3))
        self.bn = tf.keras.layers.LayerNormalization()
        self.dense = tf.keras.layers.Dense(dim,
                                           name='output_dense', 
                                           use_bias=False,
                                           kernel_regularizer=tf.keras.regularizers.l2(0.0001))
        if loss_type == 'arc':
            self.margin = ArcMarginProduct(n_cls, s=scale, m=margin, ls_eps=0.0, easy_margin=True)
        if loss_type == 'cos':
            self.margin = CosMarginProduct(n_cls, s=scale, m=margin, ls_eps=0.0)

        self.dropout = tf.keras.layers.Dropout(0.2)

    @tf.function
    def call(self, inputs, training=None, mask=None):
        x, y = inputs
        x = self.backbone(x)
        x = self.bn(x)
        x = self.dropout(x, training)
        x = self.dense(x)
        if training:
            x = self.margin([x, y])
            y_pred = tf.keras.activations.softmax(x)
            return tf.keras.losses.SparseCategoricalCrossentropy()(y, y_pred)
        else:
            return tf.math.l2_normalize(x, axis=-1)


if __name__ == '__main__':
    # arc = create_arc(input_shape=(224, 224, 3), n_cls=10, dense=128, dropout=0.2)
    # print(arc.summary())
    model = FaceModel(loss_type='cos')
    raw = tf.io.read_file("../UNet/SKU130770.png")
    image = tf.image.decode_png(raw, channels=3)
    image = tf.image.resize(images=image, size=[224, 224])
    image = image / 127.5 - 1
    image = tf.expand_dims(image, 0)
    inputs = (image, [0])
    res = model(inputs, False)
    print(res)
